package main

import (
	"context"
	"encoding/json"
	"flag"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"dns-tools/internal/handler"

	"github.com/gorilla/mux"
	"github.com/joho/godotenv"
	"github.com/sirupsen/logrus"
)

var (
	port       = flag.String("port", "", "Server port (overrides PORT env var)")
	logLevel   = flag.String("log-level", "info", "Log level (debug, info, warn, error)")
	configFile = flag.String("config", "", "Configuration file path")
)

type Config struct {
	DNSServers []string `json:"dns_servers"`
}

func main() {
	flag.Parse()

	// 加载 .env 文件（如果存在）
	// 忽略错误，因为 .env 文件可能不存在（比如在生产环境中使用系统环境变量）
	_ = godotenv.Load()

// 处理端口配置：优先使用命令行参数，其次环境变量，最后默认值
	serverPort := *port
	if serverPort == "" {
	serverPort = os.Getenv("PORT")
	if serverPort == "" {
	serverPort = "18000" // 默认端口
	}
	}
	// 初始化日志
	logger := initLogger(*logLevel)

	// 加载配置
	config := loadConfig(*configFile, logger)

	// 获取 DNS 配置
	dnsConfig := getDNSConfig(config)

	// 创建路由器
	router := mux.NewRouter()

	// 创建处理器集合
	handlers := handler.NewHandlers(logger, dnsConfig)
	batchVerifyHandler := handler.DCVHandler(logger)

	// API 路由（新版统一接口）
	api := router.PathPrefix("/api").Subrouter()
	api.HandleFunc("/health", healthHandler).Methods("GET", "HEAD", "OPTIONS")
	api.HandleFunc("/regions", regionsHandler).Methods("GET", "OPTIONS")
	api.HandleFunc("/dns/query", handler.ApplyCORS(handlers.DNSQueryHandler)).Methods("POST", "OPTIONS")
	api.HandleFunc("/domain/issue-verify", handler.ApplyCORS(handlers.IssueVerifyBatchAPIHandler)).Methods("POST", "OPTIONS")
	api.HandleFunc("/link/query", handler.ApplyCORS(handlers.LinkVerifyHandler)).Methods("POST", "OPTIONS")
	api.HandleFunc("/whois/query", handler.ApplyCORS(handlers.WhoisQueryHandler)).Methods("POST", "OPTIONS")
	api.HandleFunc("/get", handler.ApplyCORS(handlers.GetLinkHandler)).Methods("POST", "OPTIONS")
	api.HandleFunc("/dcv/verify", handler.DCVVerifyMiddleware(batchVerifyHandler.Handle)).Methods("POST", "OPTIONS")

	// 去除旧接口，仅保留 /api/**
	router.HandleFunc("/health", healthHandler).Methods("GET")

	// 静态文件服务
	router.PathPrefix("/static/").Handler(http.StripPrefix("/static/", http.FileServer(http.Dir("./static"))))
	router.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, "./static/index.html")
	})

	// 启动服务器
	addr := ":" + serverPort
	logger.WithField("port", serverPort).Info("Starting DNS tools server")

	server := &http.Server{
		Addr:         addr,
		Handler:      loggingMiddleware(logger)(router),
		ReadTimeout:  15 * time.Second,
		WriteTimeout: 15 * time.Second,
		IdleTimeout:  60 * time.Second,
	}

	// 优雅关闭
	go func() {
		sigChan := make(chan os.Signal, 1)
		signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
		<-sigChan

		logger.Info("Shutting down server...")

		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()

		if err := server.Shutdown(ctx); err != nil {
			logger.WithError(err).Error("Server shutdown error")
		} else {
			logger.Info("Server shutdown gracefully")
		}
	}()

	// 日志在 createDNSResolver 中已经输出了具体的 DNS 配置
	logger.WithField("port", serverPort).Info("Server started successfully")

	if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		logger.WithError(err).Fatal("Server failed to start")
	}
}

// getDNSConfig 获取 DNS 配置
// 优先级：环境变量 DNS_PROVIDER > DNS_SERVERS > 默认值
func getDNSConfig(config *Config) string {
	// 1. 检查 DNS_PROVIDER 环境变量（预设服务商）
	if provider := strings.TrimSpace(os.Getenv("DNS_PROVIDER")); provider != "" {
		return provider
	}

	// 2. 检查 DNS_SERVERS 环境变量（自定义服务器）
	if servers := strings.TrimSpace(os.Getenv("DNS_SERVERS")); servers != "" {
		return servers
	}

	// 3. 使用配置文件中的 DNSServers
	if len(config.DNSServers) >= 2 {
		return config.DNSServers[0] + "," + config.DNSServers[1]
	}
	if len(config.DNSServers) == 1 {
		return config.DNSServers[0]
	}

	// 4. 默认为空字符串（将使用 Google Public DNS）
	return ""
}

// loadConfig 加载配置
func loadConfig(configFile string, logger *logrus.Logger) *Config {
	config := &Config{
		DNSServers: []string{"8.8.8.8", "8.8.4.4"}, // 默认DNS服务器（用于向后兼容）
	}

	// 优先从配置文件加载
	if configFile != "" {
		data, err := os.ReadFile(configFile)
		if err == nil {
			if err := json.Unmarshal(data, config); err != nil {
				logger.WithError(err).Warn("Failed to parse config file, using defaults")
			} else {
				logger.WithField("config", configFile).Info("Loaded configuration")
			}
		} else {
			logger.WithError(err).Warn("Failed to read config file, using defaults")
		}
	}

	// 环境变量可以覆盖配置文件（优先级：环境变量 > 配置文件 > 默认值）
	if dnsServers := strings.TrimSpace(os.Getenv("DNS_SERVERS")); dnsServers != "" {
		// 按逗号分割并去除空格
		servers := strings.Split(dnsServers, ",")
		var cleanServers []string
		for _, server := range servers {
			if s := strings.TrimSpace(server); s != "" {
				cleanServers = append(cleanServers, s)
			}
		}
		if len(cleanServers) > 0 {
			config.DNSServers = cleanServers
		}
	}

	return config
}

// initLogger 初始化日志
func initLogger(level string) *logrus.Logger {
	logger := logrus.New()
	logger.SetFormatter(&logrus.JSONFormatter{
		TimestampFormat: time.RFC3339,
	})

	switch level {
	case "debug":
		logger.SetLevel(logrus.DebugLevel)
	case "warn":
		logger.SetLevel(logrus.WarnLevel)
	case "error":
		logger.SetLevel(logrus.ErrorLevel)
	default:
		logger.SetLevel(logrus.InfoLevel)
	}

	return logger
}

// loggingMiddleware 日志中间件
func loggingMiddleware(logger *logrus.Logger) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			start := time.Now()

			// 包装 ResponseWriter 以捕获状态码
			wrapped := &responseWriter{
				ResponseWriter: w,
				statusCode:     http.StatusOK,
			}

			next.ServeHTTP(wrapped, r)

			logger.WithFields(logrus.Fields{
				"method":     r.Method,
				"path":       r.URL.Path,
				"remote_ip":  r.RemoteAddr,
				"status":     wrapped.statusCode,
				"duration":   time.Since(start).Milliseconds(),
				"user_agent": r.UserAgent(),
			}).Info("Request processed")
		})
	}
}

// responseWriter 包装 ResponseWriter 以捕获状态码
type responseWriter struct {
	http.ResponseWriter
	statusCode int
}

func (w *responseWriter) WriteHeader(code int) {
	w.statusCode = code
	w.ResponseWriter.WriteHeader(code)
}

// healthHandler 健康检查处理器
func healthHandler(w http.ResponseWriter, r *http.Request) {
    data := map[string]interface{}{
        "status":    "healthy",
        "timestamp": time.Now(),
        "version":   "2.0.0",
    }
    resp := map[string]interface{}{
        "code": 1,
        "data": data,
    }
    w.Header().Set("Content-Type", "application/json")
    // 防止任何形式的缓存
    w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
    w.Header().Set("Pragma", "no-cache")
    w.Header().Set("Expires", "0")
    json.NewEncoder(w).Encode(resp)
}

// regionsHandler 返回可用区列表（从 regions.json 读取），兼容静态页区域选择
func regionsHandler(w http.ResponseWriter, r *http.Request) {
	type Region struct {
		ID      string `json:"id"`
		Name    string `json:"name"`
		URL     string `json:"url"`
		Default bool   `json:"default,omitempty"`
	}
	type RegionsConfig struct {
		Regions []Region `json:"regions"`
	}

	w.Header().Set("Content-Type", "application/json")

	// 读取本地 regions.json（若不存在则返回默认）
	data, err := os.ReadFile("regions.json")
	if err != nil {
		// 默认区域
		resp := map[string]interface{}{
			"code": 1,
			"data": map[string]interface{}{
				"regions":       []Region{{ID: "default", Name: "本地", URL: "http://localhost:18000", Default: true}},
				"currentRegion": "default",
			},
		}
		json.NewEncoder(w).Encode(resp)
		return
	}

	var cfg RegionsConfig
	if err := json.Unmarshal(data, &cfg); err != nil || len(cfg.Regions) == 0 {
		resp := map[string]interface{}{
			"code": 1,
			"data": map[string]interface{}{
				"regions":       []Region{{ID: "default", Name: "本地", URL: "http://localhost:18000", Default: true}},
				"currentRegion": "default",
			},
		}
		json.NewEncoder(w).Encode(resp)
		return
	}

	// 根据当前访问的 Host 判断实际区域
	current := ""
	host := r.Host
	
	// 获取完整的请求 URL
	scheme := "http"
	if r.TLS != nil {
		scheme = "https"
	}
	// 检查 X-Forwarded-Proto 头
	if proto := r.Header.Get("X-Forwarded-Proto"); proto != "" {
		scheme = proto
	}
	
	// 构建当前访问的完整 URL（不包含路径）
	currentURL := scheme + "://" + host
	
	// 根据当前 URL 找到对应的区域
	for _, rgn := range cfg.Regions {
		// 比较 URL 的主机部分
		if strings.HasPrefix(rgn.URL, currentURL) || strings.Contains(rgn.URL, host) {
			current = rgn.ID
			break
		}
	}
	
	// 如果没有匹配的，使用默认区域
	if current == "" {
		for _, rgn := range cfg.Regions {
			if rgn.Default {
				current = rgn.ID
				break
			}
		}
	}
	if current == "" {
		current = cfg.Regions[0].ID
	}

	resp := map[string]interface{}{
		"code": 1,
		"data": map[string]interface{}{
			"regions":       cfg.Regions,
			"currentRegion": current,
		},
	}
	w.Header().Set("Content-Type", "application/json")
	// 防止任何形式的缓存
	w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
	w.Header().Set("Pragma", "no-cache")
	w.Header().Set("Expires", "0")
	json.NewEncoder(w).Encode(resp)
}
